{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {class}`~tvm.relay.transform.MergeComposite`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "合并复合（merge composite）过程旨在将多个符合特定模式的 Relay 算子合并，并将它们组合成单一的 Relay 函数。\n",
    "\n",
    "例如，假设有如下计算图：\n",
    "\n",
    "```\n",
    "    conv2d\n",
    "      |            (合并复合过程)\n",
    "   bias_add            ====>           conv2d_bias_relu\n",
    "      |            (我们的目标)\n",
    "     relu\n",
    "```\n",
    "\n",
    "\n",
    "在合并复合过程之前的 Relay IR：\n",
    "\n",
    "```\n",
    "    fn (%data: Tensor[(1, 512, 28, 28), float32], %kernel: Tensor[(256, 512, 1, 1), float32],\n",
    "            %bias: Tensor[(256), float32]) -> Tensor[(1, 256, 28, 28), float32] {\n",
    "        %0 = nn.conv2d(%data, %kernel, kernel_size=[1, 1])\n",
    "            /* ty=Tensor[(1, 256, 28, 28), float32] */;\n",
    "        %1 = nn.bias_add(%0, %bias) /* ty=Tensor[(1, 256, 28, 28), float32] */;\n",
    "        nn.relu(%1) /* ty=Tensor[(1, 256, 28, 28), float32] */\n",
    "    }\n",
    "```\n",
    "\n",
    "在合并复合过程之后的 Relay IR：\n",
    "\n",
    "```\n",
    "    fn (%data: Tensor[(1, 512, 28, 28), float32], %kernel: Tensor[(256, 512, 1, 1), float32],\n",
    "            %bias: Tensor[(256), float32]) -> Tensor[(1, 256, 28, 28), float32] {\n",
    "      %2 = fn (%x: Tensor[(1, 512, 28, 28), float32], %y: Tensor[(256, 512, 1, 1), float32],\n",
    "            %z: Tensor[(256), float32], Primitive=1, Composite=\"conv2d_bias_relu\") ->\n",
    "            Tensor[(1, 256, 28, 28), float32] {\n",
    "        %0 = nn.conv2d(%x, %y, kernel_size=[1, 1]) /* ty=Tensor[(1, 256, 28, 28), float32] */;\n",
    "        %1 = nn.bias_add(%0, %z) /* ty=Tensor[(1, 256, 28, 28), float32] */;\n",
    "        nn.relu(%1) /* ty=Tensor[(1, 256, 28, 28), float32] */\n",
    "      };\n",
    "      %2(%data, %kernel, %bias) /* ty=Tensor[(1, 256, 28, 28), float32] */\n",
    "    }\n",
    "```\n",
    "\n",
    "\n",
    "正如你在第二个 Relay 示例中所看到的，指定的模式被封装在函数中。然后调用该函数，产生与第一个 Relay 示例相同的结果。\n",
    "\n",
    "这个合并复合过程的便利用途是将多个算子卸载到单一的外部代码生成函数中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "from tvm.relay.dataflow_pattern import TupleGetItemPattern, is_op, wildcard\n",
    "\n",
    "def make_add_sub_mul_pattern():\n",
    "    r\"\"\"Create a pattern to match the following graph.\n",
    "\n",
    "    add  sub\n",
    "     \\   /\n",
    "      \\ /\n",
    "      mul\n",
    "    \"\"\"\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    return (x + y) * (x - y)\n",
    "\n",
    "\n",
    "def make_add_relu_pattern():\n",
    "    r\"\"\"Create a pattern to match the following graph.\n",
    "\n",
    "     add\n",
    "      |\n",
    "    relu\n",
    "    \"\"\"\n",
    "    add_node = wildcard() + wildcard()\n",
    "    r = is_op(\"nn.relu\")(add_node)\n",
    "    return r\n",
    "\n",
    "\n",
    "def make_conv_bias_relu_pattern():\n",
    "    r\"\"\"Create a pattern to match the following graph.\n",
    "\n",
    "     conv2d\n",
    "       |\n",
    "    bias_add\n",
    "       |\n",
    "     relu\n",
    "    \"\"\"\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    z = wildcard()\n",
    "    conv_node = is_op(\"nn.conv2d\")(x, y)\n",
    "    bias_node = is_op(\"nn.bias_add\")(conv_node, z)\n",
    "    r = is_op(\"nn.relu\")(bias_node)\n",
    "    return r\n",
    "\n",
    "\n",
    "def make_pattern_with_optional():\n",
    "    r\"\"\"Create a pattern to match the following graph. Note that relu is optinal.\n",
    "\n",
    "     conv2d\n",
    "       |\n",
    "    bias_add\n",
    "       |\n",
    "     (relu)\n",
    "    \"\"\"\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    z = wildcard()\n",
    "    conv_node = is_op(\"nn.conv2d\")(x, y)\n",
    "    bias_node = is_op(\"nn.bias_add\")(conv_node, z)\n",
    "    r = bias_node.optional(lambda x: is_op(\"nn.relu\")(x))\n",
    "    return r\n",
    "\n",
    "\n",
    "def make_add_add_add_pattern():\n",
    "    r\"\"\"Create a pattern to match the following graph.\n",
    "       Useful for testing re-using a call node.\n",
    "\n",
    "        x    y\n",
    "      /  \\  /\n",
    "      |  add\n",
    "       \\  |  \\\n",
    "         add |\n",
    "          | /\n",
    "         add\n",
    "    \"\"\"\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    add_node = is_op(\"add\")(x, y)\n",
    "    add_node_1 = is_op(\"add\")(x, add_node)\n",
    "    r = is_op(\"add\")(add_node_1, add_node)\n",
    "    return r\n",
    "\n",
    "\n",
    "def make_bn_relu_pattern():\n",
    "    r\"\"\"Create a pattern to match the following graph.\n",
    "\n",
    "     batch_norm\n",
    "         |\n",
    "    TupleGetItem(0)\n",
    "         |\n",
    "       relu\n",
    "    \"\"\"\n",
    "    x = wildcard()\n",
    "    gamma = wildcard()\n",
    "    beta = wildcard()\n",
    "    moving_mean = wildcard()\n",
    "    moving_var = wildcard()\n",
    "    bn_node = is_op(\"nn.batch_norm\")(x, gamma, beta, moving_mean, moving_var)\n",
    "    tuple_get_item_node = TupleGetItemPattern(bn_node, 0)\n",
    "    r = is_op(\"nn.relu\")(tuple_get_item_node)\n",
    "    return r\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 简单的合并"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从简单的计算图中正确生成复合函数。期望模式 `make_add_relu_pattern` 能够被合并成单一的算子 `add_relu`。\n",
    "```\n",
    "        a  b\n",
    "        \\ /               a  b\n",
    "        add    ====>      \\ /\n",
    "         |             add_relu\n",
    "       relu\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "pattern_table = [(\"add_relu\", make_add_relu_pattern())]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "查看效果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relay, tir\n",
    "from tvm.relay.testing import run_opt_pass\n",
    "\n",
    "def before():\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "    add_node = relay.add(a, b)\n",
    "    r = relay.nn.relu(add_node)\n",
    "    return relay.Function([a, b], r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 分支合并\n",
    "\n",
    "测试从分支图中正确生成复合函数。\n",
    "\n",
    "期望模式 `make_add_sub_mul_pattern` 能够被合并成单一的算子 `add_sub_mul`。\n",
    "\n",
    "```\n",
    "       a  b  a  b\n",
    "        \\/    \\/\n",
    "        add  sub                       a  b\n",
    "         \\   /                          \\/\n",
    "          \\ /                      add_sub_mul\n",
    "          mul                     c     |\n",
    "          /  \\                     \\    |\n",
    "       c /  c |       ====>        add_sub_mul\n",
    "       \\/   \\/                          |\n",
    "       add  sub                         |\n",
    "        \\   /                         relu\n",
    "         \\ /\n",
    "         mul\n",
    "          |\n",
    "          |\n",
    "        relu\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>c: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>c, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>c, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>c: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_subtract_multiply_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_sub_mul&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_subtract_multiply_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_sub_mul&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>c, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [(\"add_sub_mul\", make_add_sub_mul_pattern())]\n",
    "def before():\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "    c = relay.var(\"c\", shape=(10, 10))\n",
    "    add_node = relay.add(a, b)\n",
    "    sub_node = relay.subtract(a, b)\n",
    "    mul_node = relay.multiply(add_node, sub_node)\n",
    "    add_node_2 = relay.add(c, mul_node)\n",
    "    sub_node_2 = relay.subtract(c, mul_node)\n",
    "    mul_node_2 = relay.multiply(add_node_2, sub_node_2)\n",
    "    r = relay.nn.relu(mul_node_2)\n",
    "    return relay.Function([a, b, c], r)\n",
    "\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 重用调用合并\n",
    "\n",
    "测试从重用调用节点的简单图中正确生成复合函数。\n",
    "\n",
    "期望模式 `make_add_add_add` 能够被合并成单一的算子 `add_add_add`。\n",
    "\n",
    "```\n",
    "        x     y\n",
    "         \\   / \\\n",
    "          sub  |           x     y\n",
    "        /  |  /             \\   / |\n",
    "        | add      ====>     sub  |\n",
    "         \\ |  \\               |  /\n",
    "          add |           add_add_add\n",
    "           | /\n",
    "          add\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_add_add_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_add_add&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "    sub_node = relay.subtract(a, b)\n",
    "\n",
    "    # pattern\n",
    "    add_node = relay.add(sub_node, b)\n",
    "    add_node_1 = relay.add(sub_node, add_node)\n",
    "    r = relay.add(add_node_1, add_node)\n",
    "\n",
    "    return relay.Function([a, b], r)\n",
    "\n",
    "\n",
    "pattern_table = [(\"add_add_add\", make_add_add_add_pattern())]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 合并多个模式\n",
    "\n",
    "测试图中不同模式是否正确合并。\n",
    "\n",
    "期望模式 `make_conv_bias_relu_pattern` 能够被合并成单一的算子 `conv_bias_relu`。同时，也期望 `make_add_relu_pattern` 能够被合并成单一的算子 `add_relu`。\n",
    "```\n",
    "        data   kernel\n",
    "          \\      /\n",
    "           \\    /\n",
    "           conv2d                   data   kernel   bias\n",
    "             |                         \\      |      /\n",
    "             |   bias                 conv2d_bias_relu\n",
    "             |   /                            |\n",
    "          bias_add        ====>               |    a\n",
    "             |                                |   /\n",
    "           relu  a                        add_relu\n",
    "             \\  /                             |\n",
    "             add                              |  b\n",
    "              |                               | /\n",
    "            relu  b                          mul\n",
    "              |  /\n",
    "             mul\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel: Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias: Tensor[(<span style=\"color: #008000\">256</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>a);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>);\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel: Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias: Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.conv2d_nn.bias_add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;conv2d_bias_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_11, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32], Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], Tensor[(<span style=\"color: #008000\">256</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32], Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>a) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    data = relay.var(\"data\", shape=(1, 512, 28, 28))\n",
    "    kernel = relay.var(\"kernel\", shape=(256, 512, 1, 1))\n",
    "    bias = relay.var(\"bias\", shape=(256,))\n",
    "    a = relay.var(\"a\", shape=(1, 256, 28, 28))\n",
    "    b = relay.var(\"b\", shape=(1, 256, 28, 28))\n",
    "\n",
    "    conv_node = relay.nn.conv2d(\n",
    "        data, kernel, kernel_size=(1, 1), padding=(0, 0), strides=(1, 1)\n",
    "    )\n",
    "\n",
    "    bias_node = relay.nn.bias_add(conv_node, bias)\n",
    "    relu_node = relay.nn.relu(bias_node)\n",
    "    add_node = relay.add(relu_node, a)\n",
    "    relu_node_2 = relay.nn.relu(add_node)\n",
    "    r = relay.multiply(relu_node_2, b)\n",
    "    return relay.Function([data, kernel, bias, a, b], r)\n",
    "\n",
    "\n",
    "pattern_table = [\n",
    "    (\"conv2d_bias_relu\", make_conv_bias_relu_pattern()),\n",
    "    (\"add_relu\", make_add_relu_pattern()),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 合并可选模式\n",
    "\n",
    "测试包含可选算子的模式。可以定义包含某些可选算子的模式。合并复合过程将为所有匹配的模式创建复合函数，但会带有不同的 \"PartitionedFromPattern\" 属性。期望后端代码生成器能够分析该属性并确定相应的算子。\n",
    "```\n",
    "模式:        匹配情况 A:        匹配情况 B:\n",
    "\n",
    " conv2d        conv2d             conv2d\n",
    "   |             |                  |\n",
    "bias_add      bias_add           bias_add\n",
    "   |             |\n",
    " (relu)         relu\n",
    "```\n",
    "在上面的示例中，匹配情况 A 的复合函数将具有 `PartitionedFromPattern=\"nn.conv2d_nn.bias_add_nn.relu_\"`，而匹配情况 B 的复合函数将是 `\"nn.conv2d_nn.bias_add_\"`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>w1: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">3</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b1: Tensor[(<span style=\"color: #008000\">3</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b1)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w1: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b1: Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_2: Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.conv2d_nn.bias_add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;layer&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32], Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], Tensor[(<span style=\"color: #008000\">3</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.conv2d_nn.bias_add_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;layer&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32], Tensor[(<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], Tensor[(<span style=\"color: #008000\">3</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(1, 3, 7, 7))\n",
    "    w1 = relay.var(\"w\", shape=(3, 3, 1, 1))\n",
    "    b1 = relay.var(\"b\", shape=(3,))\n",
    "    w2 = relay.var(\"w\", shape=(3, 3, 1, 1))\n",
    "    b2 = relay.var(\"b\", shape=(3,))\n",
    "    conv = relay.nn.conv2d(x, w1, kernel_size=(1, 1))\n",
    "    bias = relay.nn.bias_add(conv, b1)\n",
    "    relu = relay.nn.relu(bias)\n",
    "    conv = relay.nn.conv2d(relu, w2, kernel_size=(1, 1))\n",
    "    bias = relay.nn.bias_add(conv, b2)\n",
    "    return relay.Function([x, w1, w2, b1, b2], bias)\n",
    "\n",
    "\n",
    "pattern_table = [(\"layer\", make_pattern_with_optional())]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 依序合并\n",
    "\n",
    "测试模式是否按照它们在模式表中的顺序进行合并。\n",
    "\n",
    "在某些情况下，一个模式可能是另一个模式的子图，此时不清楚哪个匹配应该优先。优先级应取决于模式在模式表中声明的顺序。最先声明的模式将以最高优先级合并，而最后声明的模式将以最低优先级合并。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.dataflow_pattern import is_op, wildcard\n",
    "def pattern_A():\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    out = is_op(\"add\")(x, y)\n",
    "    out = is_op(\"abs\")(out)\n",
    "    out = is_op(\"nn.relu\")(out)\n",
    "    return out\n",
    "\n",
    "def pattern_B():\n",
    "    x = wildcard()\n",
    "    y = wildcard()\n",
    "    out = is_op(\"add\")(x, y)\n",
    "    out = is_op(\"abs\")(out)\n",
    "    return out\n",
    "\n",
    "def pattern_C():\n",
    "    x = wildcard()\n",
    "    out = is_op(\"abs\")(x)\n",
    "    out = is_op(\"nn.relu\")(out)\n",
    "    return out\n",
    "\n",
    "def before():\n",
    "    input_1 = relay.var(\"input_1\", shape=(10, 10))\n",
    "    input_2 = relay.var(\"input_2\", shape=(10, 10))\n",
    "    out = relay.add(input_1, input_2)\n",
    "    out = relay.abs(out)\n",
    "    out = relay.nn.relu(out)\n",
    "    return relay.Function([input_1, input_2], out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "检查 A 是否具有最高优先级："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_abs_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;A&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [\n",
    "    (\"A\", pattern_A()),\n",
    "    (\"B\", pattern_B()),\n",
    "    (\"C\", pattern_C()),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "检查 B 是否具有最高优先级："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_abs_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;B&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [\n",
    "    (\"B\", pattern_B()),\n",
    "    (\"C\", pattern_C()),\n",
    "    (\"A\", pattern_A()),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "检查 C 是否具有最高优先级："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;abs_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;C&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> abs(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [\n",
    "    (\"C\", pattern_C()),\n",
    "    (\"A\", pattern_A()),\n",
    "    (\"B\", pattern_B()),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 并行合并\n",
    "\n",
    "测试依赖于相同输入的并行模式是否正确合并。\n",
    "\n",
    "测试图难以用 ASCII 图形表示。它本质上是两个并行的 ``add-sub-mul`` 单元，它们都消耗 `input_1` 和 `input_2`，并将它们的结果相乘以生成输出。期望两个并行分支都能被合并，并且它们仍然应该消耗相同的输入变量 `input_1` 和 `input_2`。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>);\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_subtract_multiply_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_sub_mul&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_subtract_multiply_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_sub_mul&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    input_1 = relay.var(\"input_1\", shape=(10, 10))\n",
    "    input_2 = relay.var(\"input_2\", shape=(10, 10))\n",
    "    branch_1_add = relay.add(input_1, input_2)\n",
    "    branch_1_sub = relay.subtract(input_1, input_2)\n",
    "    branch_1 = relay.multiply(branch_1_add, branch_1_sub)\n",
    "    branch_2_add = relay.add(input_1, input_2)\n",
    "    branch_2_sub = relay.subtract(input_1, input_2)\n",
    "    branch_2 = relay.multiply(branch_2_add, branch_2_sub)\n",
    "    out = relay.multiply(branch_1, branch_2)\n",
    "    return relay.Function([input_1, input_2], out)\n",
    "\n",
    "pattern_table = [(\"add_sub_mul\", make_add_sub_mul_pattern())]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 多输入子图\n",
    "\n",
    "测试多个输入子图连接到另一个子图的情况。\n",
    "\n",
    "````\n",
    "     (1)    (2)    (3)    (4)\n",
    "    add    add    add    add\n",
    "     |      |      |      |\n",
    "    relu   relu   relu   relu\n",
    "     \\      /      \\      /\n",
    "      \\   /         \\   /\n",
    "       add           sub\n",
    "        \\            /\n",
    "          \\        /\n",
    "            \\    /\n",
    "              mul\n",
    "\n",
    "    ----> 当 1=3 且 2=4 时（情况 'A'）\n",
    "\n",
    "    add_relu  add_relu\n",
    "       \\         /\n",
    "        \\      /\n",
    "       add_sub_mul\n",
    "\n",
    "    ----> 当 1!=3 且 2!=4 时（情况 'B'）\n",
    "\n",
    "    add_relu  add_relu  add_relu  add_relu\n",
    "       \\       /           \\       /\n",
    "         \\   /               \\   /\n",
    "          add                 sub\n",
    "           \\                  /\n",
    "            --------     -----\n",
    "                   \\    /\n",
    "                    mul\n",
    "````\n",
    "\n",
    "行为上的差异源于 `add_sub_mul` 期望 `add` 和 `sub` 的输入是相同的（相同的两个 `relay` 表达式）。因此，当你有 4 个独立的输入时，模式不应被合并。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>);\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_11: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_11) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_subtract_multiply_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_sub_mul&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    before_funcs = {}\n",
    "    inputs = [relay.var(\"input_\" + str(i), shape=(10, 10)) for i in range(8)]\n",
    "    add_relu_1 = relay.add(inputs[0], inputs[1])\n",
    "    add_relu_1 = relay.nn.relu(add_relu_1)\n",
    "    add_relu_2 = relay.add(inputs[2], inputs[3])\n",
    "    add_relu_2 = relay.nn.relu(add_relu_2)\n",
    "    add_relu_3 = relay.add(inputs[4], inputs[5])\n",
    "    add_relu_3 = relay.nn.relu(add_relu_3)\n",
    "    add_relu_4 = relay.add(inputs[6], inputs[7])\n",
    "    add_relu_4 = relay.nn.relu(add_relu_4)\n",
    "    add = relay.add(add_relu_1, add_relu_2)\n",
    "    sub = relay.subtract(add_relu_3, add_relu_4)\n",
    "    out = relay.multiply(add, sub)\n",
    "    before_funcs[\"B\"] = relay.Function(inputs, out)\n",
    "    sub = relay.subtract(add_relu_1, add_relu_2)\n",
    "    out = relay.multiply(add, sub)\n",
    "    before_funcs[\"A\"] = relay.Function(inputs[:4], out)\n",
    "    return before_funcs\n",
    "\n",
    "pattern_table = [\n",
    "    (\"add_sub_mul\", make_add_sub_mul_pattern()),\n",
    "    (\"add_relu\", make_add_relu_pattern()),\n",
    "]\n",
    "graph = before()[\"A\"]\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_4: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_5: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_6: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_7: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_4, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_5);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_6, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_7);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>);\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_4: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_5: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_6: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_7: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_3_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_3_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_3_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_3_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_2_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_2_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_2_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_2_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_3) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_1_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_4, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_5) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input_6, <span style=\"color: #AA22FF; font-weight: bold\">%</span>input_7) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> subtract(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "graph = before()[\"B\"]\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 合并 `TupleGetItem`\n",
    "测试可以从包含 `TupleGetItem` 节点的模式中合并复合函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>gamma: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>beta: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_mean: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_var: Tensor[(<span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_var);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>gamma: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>beta: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_mean: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_var: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_3: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_4: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.batch_norm_TupleGetItem0_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;bn_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_3, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_4) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>gamma, <span style=\"color: #AA22FF; font-weight: bold\">%</span>beta, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_mean, <span style=\"color: #AA22FF; font-weight: bold\">%</span>moving_var) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [(\"bn_relu\", make_bn_relu_pattern())]\n",
    "\n",
    "def before():\n",
    "    x = relay.var(\"x\", shape=(1, 8))\n",
    "    gamma = relay.var(\"gamma\", shape=(8,))\n",
    "    beta = relay.var(\"beta\", shape=(8,))\n",
    "    moving_mean = relay.var(\"moving_mean\", shape=(8,))\n",
    "    moving_var = relay.var(\"moving_var\", shape=(8,))\n",
    "    bn_node = relay.nn.batch_norm(x, gamma, beta, moving_mean, moving_var)\n",
    "    tuple_get_item_node = bn_node[0]\n",
    "    r = relay.nn.relu(tuple_get_item_node)\n",
    "    return relay.Function([x, gamma, beta, moving_mean, moving_var], r)\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ``with`` 检查"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(1, 10, 10, 10))\n",
    "    w = relay.var(\"w\", shape=(10, 10, 3, 3))\n",
    "    b = relay.var(\"b\", shape=(8,))\n",
    "    conv = relay.nn.conv2d(x, w, kernel_size=(3, 3), kernel_layout=\"OIHW\", data_layout=\"NHWC\")\n",
    "    bias = relay.nn.bias_add(conv, b)\n",
    "    relu = relay.nn.relu(bias)\n",
    "    return relay.Function([x, w, b], relu)\n",
    "\n",
    "def _check_true(extract):\n",
    "    conv = extract.args[0].args[0]\n",
    "    return conv.attrs.data_layout == \"NHWC\"\n",
    "\n",
    "def _check_false(extract):\n",
    "    conv = extract.args[0].args[0]\n",
    "    return conv.attrs.data_layout == \"NCHW\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [(\"conv_bias_relu\", make_conv_bias_relu_pattern(), _check_false)]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b);\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.conv2d_nn.bias_add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;conv_bias_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [(\"conv_bias_relu\", make_conv_bias_relu_pattern(), _check_true)]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## diamond 不合并的情况\n",
    "\n",
    "左侧的模式不应匹配右侧的结构。\n",
    "\n",
    "```\n",
    "    relu             relu\n",
    "     | \\              | \\\n",
    "     | clip           | add\n",
    "     |  /             |  |\n",
    "     mul              | clip\n",
    "                      |  /\n",
    "                      mul\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel: Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias: Tensor[(<span style=\"color: #008000\">256</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>f);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">255</span>f);\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel: Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">512</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias: Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">256</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #AA22FF; font-weight: bold\">%</span>kernel, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>bias) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">1</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">255</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">256</span>, <span style=\"color: #008000\">28</span>, <span style=\"color: #008000\">28</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_pattern():\n",
    "    conv = make_conv_bias_relu_pattern()\n",
    "    clip = is_op(\"clip\")(conv, wildcard(), wildcard())\n",
    "    return is_op(\"multiply\")(conv, clip)\n",
    "\n",
    "def get_net():\n",
    "    data = relay.var(\"data\", shape=(1, 512, 28, 28))\n",
    "    kernel = relay.var(\"kernel\", shape=(256, 512, 1, 1))\n",
    "    conv = relay.nn.conv2d(data, kernel, kernel_size=(1, 1), padding=(0, 0), strides=(1, 1))\n",
    "    bias = relay.nn.bias_add(conv, relay.var(\"bias\", shape=(256,)))\n",
    "    relu = relay.nn.relu(bias)\n",
    "    add = relay.op.add(relu, relay.const(1.0))\n",
    "    clip2 = relay.op.clip(add, 0, 255)\n",
    "    mul = relay.op.multiply(relu, clip2)\n",
    "    return relay.Function(relay.analysis.free_vars(mul), mul)\n",
    "\n",
    "pattern_table = [(\"pat\", get_pattern())]\n",
    "graph = get_net()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 查询张量类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(1, 10, 10, 10))\n",
    "    w = relay.var(\"w\", shape=(10, 10, 3, 3))\n",
    "    b = relay.var(\"b\", shape=(8,))\n",
    "    add = relay.op.add(x, x)\n",
    "    relu = relay.nn.relu(add)\n",
    "    conv = relay.nn.conv2d(\n",
    "        relu, w, kernel_size=(3, 3), kernel_layout=\"OIHW\", data_layout=\"NHWC\"\n",
    "    )\n",
    "    bias = relay.nn.bias_add(conv, b)\n",
    "    relu2 = relay.nn.relu(bias)\n",
    "    return run_opt_pass(relay.Function([x, w, b], relu2), relay.transform.InferType())\n",
    "\n",
    "def _check_type_true(extract):\n",
    "    conv = extract.args[0].args[0]\n",
    "    typ = conv.checked_type\n",
    "    return bool(typ.shape[0] == 1)\n",
    "\n",
    "def _check_type_false(extract):\n",
    "    conv = extract.args[0].args[0]\n",
    "    typ = conv.checked_type\n",
    "    return bool(typ.shape[0] != 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>x) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [\n",
    "    (\"add_relu\", make_add_relu_pattern()),\n",
    "    (\"conv_bias_relu\", make_conv_bias_relu_pattern(), _check_type_false),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, <span style=\"color: #AA22FF; font-weight: bold\">%</span>x) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;add_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_01) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;nn.conv2d_nn.bias_add_nn.relu_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;conv_bias_relu&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>], data_layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias_add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_2) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">3</span>), float32], Tensor[(<span style=\"color: #008000\">8</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>w, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pattern_table = [\n",
    "    (\"add_relu\", make_add_relu_pattern()),\n",
    "    (\"conv_bias_relu\", make_conv_bias_relu_pattern(), _check_type_true),\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 不会因 `einsum` 算子而导致错误"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>]);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>b, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>]);\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>);\n",
       "  einsum(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, equation<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;...ab,...cb-&gt;...ac&quot;</span>)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">20</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, PartitionedFromPattern<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;reshape_reshape_Tuple_einsum_&quot;</span>, Composite<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;einsum_reshape&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">20</span>), float32] {\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_0, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>FunctionVar_0_1, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>), float32], Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">5</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "    einsum(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, equation<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;...ab,...cb-&gt;...ac&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">20</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">20</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>a, <span style=\"color: #AA22FF; font-weight: bold\">%</span>b) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">20</span>, <span style=\"color: #008000\">20</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tvm.relay.dataflow_pattern import TuplePattern\n",
    "\n",
    "def make_einsum_reshape_pattern():\n",
    "    x = wildcard()\n",
    "    x = is_op(\"reshape\")(x) | x\n",
    "    y = wildcard()\n",
    "    y = is_op(\"reshape\")(y) | y\n",
    "    z = is_op(\"einsum\")(TuplePattern([x, y]))\n",
    "    r = is_op(\"reshape\")(z) | z\n",
    "    return r\n",
    "\n",
    "def before():\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "    c = relay.reshape(a, [20, 5])\n",
    "    d = relay.reshape(b, [20, 5])\n",
    "    r = relay.einsum([c, d], \"...ab,...cb->...ac\")\n",
    "    return relay.Function([a, b], r)\n",
    "\n",
    "pattern_table = [\n",
    "    (\n",
    "        \"einsum_reshape\",\n",
    "        make_einsum_reshape_pattern(),\n",
    "    )\n",
    "]\n",
    "graph = before()\n",
    "tvm.IRModule.from_expr(graph).show()\n",
    "result = run_opt_pass(\n",
    "    graph, relay.transform.MergeComposite(pattern_table), import_prelude=False\n",
    ")\n",
    "assert not relay.analysis.free_vars(result), f\"在{result}图中发现了自由变量\"\n",
    "tvm.IRModule.from_expr(result).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
